#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Word 文档排版 AI Agent
一个能够根据用户提供的排版要求，将 Markdown 格式的文本内容自动排版并生成 Word 文档的程序。
"""

import re
import markdown
from docx import Document
from docx.shared import Pt, Cm, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.enum.style import WD_STYLE_TYPE
from docx.oxml.ns import qn
from docx.oxml.shared import OxmlElement, qn
from bs4 import BeautifulSoup
import argparse
import sys
import os


class FormatRuleParser:
    """排版规则解析器"""
    
    def __init__(self):
        # 中文字号映射
        self.chinese_font_sizes = {
            '初号': 42, '小初': 36,
            '一号': 26, '小一': 24,
            '二号': 22, '小二': 18,
            '三号': 16, '小三': 15,
            '四号': 14, '小四': 12,
            '五号': 10.5, '小五': 9,
            '六号': 7.5, '小六': 6.5,
            '七号': 5.5, '八号': 5
        }
        
        # 对齐方式映射
        self.alignment_map = {
            '左对齐': WD_ALIGN_PARAGRAPH.LEFT,
            '居中': WD_ALIGN_PARAGRAPH.CENTER,
            '右对齐': WD_ALIGN_PARAGRAPH.RIGHT,
            '两端对齐': WD_ALIGN_PARAGRAPH.JUSTIFY,
            '分散对齐': WD_ALIGN_PARAGRAPH.DISTRIBUTE
        }
    
    def parse_font_size(self, size_str):
        """解析字号"""
        size_str = size_str.strip()
        
        # 中文字号
        if size_str in self.chinese_font_sizes:
            return Pt(self.chinese_font_sizes[size_str])
        
        # 数字字号
        if size_str.replace('.', '').isdigit():
            return Pt(float(size_str))
        
        # 带单位的字号
        if '磅' in size_str:
            num = re.findall(r'[\d.]+', size_str)[0]
            return Pt(float(num))
        
        return Pt(12)  # 默认值
    
    def parse_spacing(self, spacing_str, base_font_size=12):
        """解析间距"""
        spacing_str = spacing_str.strip()
        
        if '行' in spacing_str:
            lines = float(re.findall(r'[\d.]+', spacing_str)[0])
            return Pt(lines * base_font_size * 1.2)  # 行间距约为字号的1.2倍
        
        if '磅' in spacing_str:
            points = float(re.findall(r'[\d.]+', spacing_str)[0])
            return Pt(points)
        
        if 'cm' in spacing_str.lower():
            cm = float(re.findall(r'[\d.]+', spacing_str)[0])
            return Cm(cm)
        
        return Pt(0)
    
    def parse_indent(self, indent_str, base_font_size=12):
        """解析缩进"""
        indent_str = indent_str.strip()
        
        if '字符' in indent_str:
            chars = float(re.findall(r'[\d.]+', indent_str)[0])
            return Pt(chars * base_font_size * 0.5)  # 中文字符宽度约为字号的0.5倍
        
        if '磅' in indent_str:
            points = float(re.findall(r'[\d.]+', indent_str)[0])
            return Pt(points)
        
        if 'cm' in indent_str.lower():
            cm = float(re.findall(r'[\d.]+', indent_str)[0])
            return Cm(cm)
        
        return Pt(0)
    
    def parse_line_spacing(self, spacing_str):
        """解析行距"""
        spacing_str = spacing_str.strip()
        
        if '倍' in spacing_str:
            multiplier = float(re.findall(r'[\d.]+', spacing_str)[0])
            return multiplier
        
        if '磅' in spacing_str:
            points = float(re.findall(r'[\d.]+', spacing_str)[0])
            return points
        
        return 1.0  # 默认单倍行距
    
    def parse_alignment(self, align_str):
        """解析对齐方式"""
        align_str = align_str.strip()
        return self.alignment_map.get(align_str, WD_ALIGN_PARAGRAPH.LEFT)
    
    def parse_margin(self, margin_str):
        """解析页边距"""
        margin_str = margin_str.strip()
        
        if 'cm' in margin_str.lower():
            cm = float(re.findall(r'[\d.]+', margin_str)[0])
            return Cm(cm)
        
        if '磅' in margin_str:
            points = float(re.findall(r'[\d.]+', margin_str)[0])
            return Pt(points)
        
        return Cm(2.5)  # 默认值
    
    def parse_format_rules(self, rules_text):
        """解析排版规则文本"""
        rules = {}
        current_style = None
        
        lines = rules_text.strip().split('\n')
        for line in lines:
            line = line.strip()
            if not line:
                continue
            
            # 识别样式名称
            if '：' in line and ('标题' in line or '正文' in line or '页边距' in line):
                parts = line.split('：', 1)
                current_style = parts[0].strip()
                if current_style not in rules:
                    rules[current_style] = {}
                
                # 解析该行的其余部分
                if len(parts) > 1:
                    self._parse_style_properties(rules[current_style], parts[1])
            
            # 继续解析当前样式的属性
            elif current_style and ('：' in line or '，' in line):
                self._parse_style_properties(rules[current_style], line)        
        return rules
    
    def _parse_style_properties(self, style_dict, properties_text):
        """解析样式属性"""
        # 将文本按逗号分割，处理各个属性
        properties = re.split(r'[，,]', properties_text)
        
        for prop in properties:
            prop = prop.strip()
            if not prop:
                continue
            
            # 使用更灵活的解析方式
            if '字体' in prop:
                if '：' in prop:
                    font_name = prop.split('：')[1].strip()
                    style_dict['font_name'] = font_name
                else:
                    # 处理 "字体 宋体" 这种格式
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['font_name'] = parts[1].strip()
            
            elif '字号' in prop:
                if '：' in prop:
                    font_size = prop.split('：')[1].strip()
                    style_dict['font_size'] = font_size
                else:
                    # 处理 "字号 小四" 这种格式
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['font_size'] = parts[1].strip()
            
            elif '对齐' in prop:
                if '：' in prop:
                    alignment = prop.split('：')[1].strip()
                    style_dict['alignment'] = alignment
                else:
                    # 处理 "对齐方式 居中" 这种格式
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['alignment'] = parts[-1].strip()
            
            elif '加粗' in prop:
                if '：' in prop:
                    bold = prop.split('：')[1].strip()
                    style_dict['bold'] = bold in ['是', 'True', 'true', '1']
                else:
                    # 处理 "加粗 是" 这种格式
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['bold'] = parts[1].strip() in ['是', 'True', 'true', '1']
            
            elif '倾斜' in prop:
                if '：' in prop:
                    italic = prop.split('：')[1].strip()
                    style_dict['italic'] = italic in ['是', 'True', 'true', '1']
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['italic'] = parts[1].strip() in ['是', 'True', 'true', '1']
            
            elif '首行缩进' in prop:
                if '：' in prop:
                    indent = prop.split('：')[1].strip()
                    style_dict['first_line_indent'] = indent
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['first_line_indent'] = ' '.join(parts[1:]).strip()
            
            elif '行距' in prop:
                if '：' in prop:
                    line_spacing = prop.split('：')[1].strip()
                    style_dict['line_spacing'] = line_spacing
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['line_spacing'] = ' '.join(parts[1:]).strip()
            
            elif '段前间距' in prop:
                if '：' in prop:
                    space_before = prop.split('：')[1].strip()
                    style_dict['space_before'] = space_before
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['space_before'] = ' '.join(parts[1:]).strip()
            
            elif '段后间距' in prop:
                if '：' in prop:
                    space_after = prop.split('：')[1].strip()
                    style_dict['space_after'] = space_after
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['space_after'] = ' '.join(parts[1:]).strip()
            
            elif '上' in prop:
                if '：' in prop:
                    top_margin = prop.split('：')[1].strip()
                    style_dict['top_margin'] = top_margin
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['top_margin'] = parts[1].strip()
            
            elif '下' in prop:
                if '：' in prop:
                    bottom_margin = prop.split('：')[1].strip()
                    style_dict['bottom_margin'] = bottom_margin
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['bottom_margin'] = parts[1].strip()
            
            elif '左' in prop:
                if '：' in prop:
                    left_margin = prop.split('：')[1].strip()
                    style_dict['left_margin'] = left_margin
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['left_margin'] = parts[1].strip()
            
            elif '右' in prop:
                if '：' in prop:
                    right_margin = prop.split('：')[1].strip()
                    style_dict['right_margin'] = right_margin
                else:
                    parts = prop.split()
                    if len(parts) >= 2:
                        style_dict['right_margin'] = parts[1].strip()


class MarkdownProcessor:
    """Markdown 处理器"""
    
    def __init__(self):
        self.md = markdown.Markdown(extensions=['extra'])
    
    def parse_markdown(self, md_text):
        """解析 Markdown 文本"""
        html = self.md.convert(md_text)
        soup = BeautifulSoup(html, 'html.parser')
        
        elements = []
        for tag in soup.find_all(['h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p', 'ul', 'ol']):
            if tag.name.startswith('h'):
                level = int(tag.name[1])
                elements.append({
                    'type': 'heading',
                    'level': level,
                    'text': tag.get_text().strip()
                })
            elif tag.name == 'p':
                text = tag.get_text().strip()
                if text:
                    elements.append({
                        'type': 'paragraph',
                        'text': text
                    })
            elif tag.name in ['ul', 'ol']:
                items = [li.get_text().strip() for li in tag.find_all('li')]
                elements.append({
                    'type': 'list',
                    'ordered': tag.name == 'ol',
                    'items': items
                })
        
        return elements


class WordDocumentGenerator:
    """Word 文档生成器"""
    
    def __init__(self):
        self.parser = FormatRuleParser()
    
    def create_document(self, format_rules, markdown_elements, output_path):
        """创建 Word 文档"""
        doc = Document()
        
        # 设置页边距
        if '页边距' in format_rules:
            self._set_page_margins(doc, format_rules['页边距'])
        
        # 创建样式
        styles = self._create_styles(doc, format_rules)
        
        # 处理 Markdown 元素
        for element in markdown_elements:
            self._add_element_to_document(doc, element, styles)
        
        # 保存文档
        doc.save(output_path)
        print(f"文档已保存到: {output_path}")
    
    def _set_page_margins(self, doc, margin_rules):
        """设置页边距"""
        section = doc.sections[0]
        
        if 'top_margin' in margin_rules:
            section.top_margin = self.parser.parse_margin(margin_rules['top_margin'])
        if 'bottom_margin' in margin_rules:
            section.bottom_margin = self.parser.parse_margin(margin_rules['bottom_margin'])
        if 'left_margin' in margin_rules:
            section.left_margin = self.parser.parse_margin(margin_rules['left_margin'])
        if 'right_margin' in margin_rules:
            section.right_margin = self.parser.parse_margin(margin_rules['right_margin'])
    
    def _create_styles(self, doc, format_rules):
        """创建样式"""
        styles = {}
        
        for style_name, style_props in format_rules.items():
            if style_name == '页边距':
                continue
            
            # 获取或创建样式
            try:
                style = doc.styles[style_name]
            except KeyError:
                style = doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
              # 设置字体
            if 'font_name' in style_props:
                style.font.name = style_props['font_name']
                # 设置中文字体
                r = style.element.rPr
                if r is None:
                    r = style.element.get_or_add_rPr()
                r.rFonts.set(qn('w:eastAsia'), style_props['font_name'])
            
            if 'font_size' in style_props:
                style.font.size = self.parser.parse_font_size(style_props['font_size'])
            
            if 'bold' in style_props:
                style.font.bold = style_props['bold']
            
            if 'italic' in style_props:
                style.font.italic = style_props['italic']
            
            # 设置段落格式
            if 'alignment' in style_props:
                style.paragraph_format.alignment = self.parser.parse_alignment(style_props['alignment'])
            
            if 'first_line_indent' in style_props:
                base_size = 12
                if 'font_size' in style_props:
                    base_size = float(re.findall(r'[\d.]+', style_props['font_size'])[0]) if re.findall(r'[\d.]+', style_props['font_size']) else 12
                style.paragraph_format.first_line_indent = self.parser.parse_indent(style_props['first_line_indent'], base_size)
            
            if 'line_spacing' in style_props:
                line_spacing = self.parser.parse_line_spacing(style_props['line_spacing'])
                if isinstance(line_spacing, float) and line_spacing > 0 and line_spacing < 10:
                    style.paragraph_format.line_spacing = line_spacing
            
            if 'space_before' in style_props:
                style.paragraph_format.space_before = self.parser.parse_spacing(style_props['space_before'])
            
            if 'space_after' in style_props:
                style.paragraph_format.space_after = self.parser.parse_spacing(style_props['space_after'])
            
            styles[style_name] = style
        
        return styles
    
    def _add_element_to_document(self, doc, element, styles):
        """向文档添加元素"""
        if element['type'] == 'heading':
            level = element['level']
            style_name = f"标题{['一', '二', '三', '四', '五', '六'][level-1]}" if level <= 6 else "标题六"
            
            paragraph = doc.add_paragraph(element['text'])
            if style_name in styles:
                paragraph.style = styles[style_name]
            else:
                # 使用内置标题样式
                paragraph.style = doc.styles[f'Heading {level}'] if level <= 9 else doc.styles['Heading 9']
        
        elif element['type'] == 'paragraph':
            paragraph = doc.add_paragraph(element['text'])
            if '正文' in styles:
                paragraph.style = styles['正文']
        
        elif element['type'] == 'list':
            for item in element['items']:
                paragraph = doc.add_paragraph(item, style='List Bullet' if not element['ordered'] else 'List Number')
                if '正文' in styles:
                    # 应用正文样式的字体设置
                    run = paragraph.runs[0] if paragraph.runs else paragraph.add_run()
                    if hasattr(styles['正文'].font, 'name') and styles['正文'].font.name:
                        run.font.name = styles['正文'].font.name
                    if hasattr(styles['正文'].font, 'size') and styles['正文'].font.size:
                        run.font.size = styles['正文'].font.size


class WordAIAgent:
    """Word 文档排版 AI Agent"""
    
    def __init__(self):
        self.rule_parser = FormatRuleParser()
        self.md_processor = MarkdownProcessor()
        self.doc_generator = WordDocumentGenerator()
    
    def process(self, format_rules_text, markdown_text, output_path):
        """处理排版任务"""
        print("解析排版规则...")
        format_rules = self.rule_parser.parse_format_rules(format_rules_text)
        
        print("解析 Markdown 内容...")
        markdown_elements = self.md_processor.parse_markdown(markdown_text)
        
        print("生成 Word 文档...")
        self.doc_generator.create_document(format_rules, markdown_elements, output_path)
        
        return True


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='Word 文档排版 AI Agent')
    parser.add_argument('--rules', '-r', required=True, help='排版规则文件路径')
    parser.add_argument('--input', '-i', required=True, help='输入的 Markdown 文件路径')
    parser.add_argument('--output', '-o', required=True, help='输出的 Word 文档路径')
    
    args = parser.parse_args()
    
    # 检查输入文件
    if not os.path.exists(args.rules):
        print(f"错误：排版规则文件不存在: {args.rules}")
        return 1
    
    if not os.path.exists(args.input):
        print(f"错误：输入 Markdown 文件不存在: {args.input}")
        return 1
    
    # 读取文件
    try:
        with open(args.rules, 'r', encoding='utf-8') as f:
            rules_text = f.read()
        
        with open(args.input, 'r', encoding='utf-8') as f:
            markdown_text = f.read()
    except Exception as e:
        print(f"错误：读取文件时发生错误: {e}")
        return 1
    
    # 创建 AI Agent 并处理
    agent = WordAIAgent()
    try:
        agent.process(rules_text, markdown_text, args.output)
        print("处理完成！")
        return 0
    except Exception as e:
        print(f"错误：处理过程中发生错误: {e}")
        return 1


if __name__ == "__main__":
    # 如果没有命令行参数，运行演示
    if len(sys.argv) == 1:
        print("Word 文档排版 AI Agent - 演示模式")
        print("=" * 50)
        
        # 示例排版规则
        sample_rules = """标题一：字体 宋体，字号 小一，对齐方式 居中，段前间距 1行，段后间距 0.5行
标题二：字体 黑体，字号 三号，加粗 是，段前间距 12磅
正文：字体 宋体，字号 小四，首行缩进 2字符，行距 1.25倍，段前间距 0磅，段后间距 0磅
页边距：上 2.5cm，下 2.5cm，左 3cm，右 3cm"""
        
        # 示例 Markdown 内容
        sample_markdown = """# 第一章 引言

这是第一章的内容。本章主要介绍了研究背景和意义。

## 1.1 研究背景

在当今信息化时代，文档处理的自动化需求日益增长。传统的手工排版方式效率低下，难以满足现代办公的需求。

## 1.2 研究意义

本研究旨在开发一个智能的文档排版系统，能够自动化处理文档格式，提高工作效率。

# 第二章 相关技术

本章介绍相关的技术背景。

## 2.1 Python-docx 库

Python-docx 是一个强大的 Python 库，用于创建和修改 Word 文档。

## 2.2 Markdown 格式

Markdown 是一种轻量级标记语言，具有简洁的语法和良好的可读性。"""
        
        # 创建演示文件
        os.makedirs('demo', exist_ok=True)
        
        with open('demo/rules.txt', 'w', encoding='utf-8') as f:
            f.write(sample_rules)
        
        with open('demo/input.md', 'w', encoding='utf-8') as f:
            f.write(sample_markdown)
        
        # 运行演示
        agent = WordAIAgent()
        try:
            agent.process(sample_rules, sample_markdown, 'demo/output.docx')
            print("\n演示完成！生成的文件：")
            print("- 排版规则: demo/rules.txt")
            print("- 输入内容: demo/input.md")
            print("- 输出文档: demo/output.docx")
        except Exception as e:
            print(f"演示过程中发生错误: {e}")
    else:
        sys.exit(main())
